local NeuralNetwork = require("NeuralNetwork");

-- 生成随机训练数据和测试数据
local function generate_data(num_samples)
    local inputs = {}
    local targets = {}
    for i = 1, num_samples do
        -- 随机生成输入值 x1, x2, x3 在 [0, 1] 范围内
        local x1 = math.random()
        local x2 = math.random()
        local x3 = math.random()
        -- 目标值为 y = x1 * 10 + x2^3 + sin(x3)
        local target = x1 * 5 + 10 * x2^3 + 5 * math.sin(x3)
        table.insert(inputs, {x1, x2, x3})
        table.insert(targets, target)
    end
    return inputs, targets
end

-- 主程序
local function main()
    -- 生成 100 条随机数据
    local num_samples = 100
    local inputs, targets = generate_data(num_samples)

    -- 将数据分为训练集和测试集（80% 训练，20% 测试）
    local train_inputs = {}
    local train_targets = {}
    local test_inputs = {}
    local test_targets = {}
    for i = 1, num_samples do
        if i <= num_samples * 0.8 then
            table.insert(train_inputs, inputs[i])
            table.insert(train_targets, targets[i])
        else
            table.insert(test_inputs, inputs[i])
            table.insert(test_targets, targets[i])
        end
    end

    -- 创建神经网络
    local input_size = 3
    local hidden_size1 = 6  -- 第一隐藏层神经元数量
    local hidden_size2 = 5  -- 第二隐藏层神经元数量
    local output_size = 1
    local epochs = 16000
    local learning_rate = 0.005
    local network = NeuralNetwork:new(input_size, hidden_size1, hidden_size2, output_size)

    -- 训练神经网络
    network:train(train_inputs, train_targets, epochs, learning_rate)

    -- 测试神经网络
    print("\nTesting the neural network:")
    network:test(test_inputs, test_targets)
end

-- 运行主程序
main()